import torch
import numpy as np
import matplotlib.pyplot as plt
from data.dataset import get_data
from model.model import Transfollower, lstm_model, nn_model
from config import Settings, HighDSettings

def test_model():
    # 加载配置
    DATASET = 'highD'
    if DATASET == 'highD':
        settings = HighDSettings()
    else:
        settings = Settings()

    MODEL = 'transfollower'
    SEQ_LEN = settings.SEQ_LEN
    LABEL_LEN = settings.LABEL_LEN
    PRED_LEN = settings.PRED_LEN
    BATCH_SIZE = settings.BATCH_SIZE
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # 加载模型
    model_path = 'checkpoints/highD_transfollower_model.pt'
    model = torch.load(model_path, map_location=device)
    model.to(device)
    model.eval()

    # 加载测试数据
    _, _, test_loader = get_data(data_name=DATASET, config=settings)

    # 定义损失函数
    criterion = torch.nn.MSELoss()

    # 预测和对比
    pred_lat_speeds = []
    true_lat_speeds = []
    test_loss = 0.0

    with torch.no_grad():
        for i, item in enumerate(test_loader):
            enc_inp = item['his'].float().to(device)  # (256,50,9)

            batch_y = item['svSpd'].float()  # (256,149,1)
            # y_label = batch_y[:,-PRED_LEN:,:].to(device)    #(256,137,1)
            batch_y_mark = item['lvSpd'].float().to(device)  # (256,149,1)

            batch_x = item['latSpd'].float()  # (256,149,1)
            x_label = batch_x[:, -PRED_LEN:, :].to(device)  # (256,137,1)
            batch_x_rd_mark = item['right_dist'].float().to(device)  # (256,149,1)
            batch_x_ld_mark = item['left_dist'].float().to(device)  # (256,149,1)
            batch_x_rs_mark = item['right_lat_spd'].float().to(device)  # (256,149,1)
            batch_x_ls_mark = item['left_lat_spd'].float().to(device)  # (256,149,1)
            batch_x_offset = item['offset'].float().to(device)  # (256,149,1)
            batch_x_l_mask = item['mask-l'].float().to(device)  # (256,149,1)
            batch_x_r_mask = item['mask-r'].float().to(device)  # (256,149,1)

            # decoder input
            sev_inp = torch.zeros([batch_x.shape[0], PRED_LEN, batch_x.shape[-1]]).float() + \
                      batch_x[:, :LABEL_LEN, :].mean(axis=1, keepdim=True)
            sev_inp = torch.cat([batch_x[:, :LABEL_LEN, :], sev_inp], dim=1).float().to(device)
            sev_inp = torch.cat(
                [sev_inp, batch_x_rd_mark, batch_x_ld_mark, batch_x_rs_mark, batch_x_ls_mark, batch_x_offset,
                 batch_x_l_mask, batch_x_r_mask], dim=-1)

            dec_inp = torch.zeros([batch_y.shape[0], PRED_LEN, batch_y.shape[-1]]).float() + \
                      batch_y[:, :LABEL_LEN, :].mean(axis=1,
                                                     keepdim=True)  # 前一个的形状是[256,137,1]填充的后一个的形状是[256,1,1],不同批次不同   自我车辆
            dec_inp = torch.cat([batch_y[:, :LABEL_LEN, :], dec_inp], dim=1).float().to(
                device)  # 拼接[256,149,1],上面12行是自车一个标签的 历史  后面137行是取得平均值
            dec_inp = torch.cat([dec_inp, batch_y_mark],
                                dim=-1)  # adding lv speed 拼接成[256,149,2]  第一列是预测纵向  第二列是前车速度  第三列是预测横向  第四五六七列是周围
            dec_inp = torch.cat([dec_inp, sev_inp], dim=-1)

            # 前向传播
            if MODEL == 'transfollower':
                out = model(enc_inp, dec_inp)[0]
            else:
                out = model(enc_inp, dec_inp)

            # 获取预测的横向速度
            pred_lat_speed = out[:, :, 0].unsqueeze(-1)
            lat_speed_loss = criterion(pred_lat_speed, x_label)

            # 记录真实和预测的横向速度
            pred_lat_speeds.extend(pred_lat_speed.cpu().numpy())
            true_lat_speeds.extend(x_label.cpu().numpy())

            test_loss += lat_speed_loss.item()

    # 计算平均损失
    test_loss /= len(test_loader)
    print(f'Test Loss: {test_loss:.7f}')

    # 转换为 numpy 数组
    pred_lat_speeds = np.concatenate(pred_lat_speeds, axis=0)
    true_lat_speeds = np.concatenate(true_lat_speeds, axis=0)

    # 绘制对比图
    plt.figure(figsize=(10, 5))
    plt.plot(pred_lat_speeds.flatten(), label="Predicted Lateral Speed")
    plt.plot(true_lat_speeds.flatten(), label="True Lateral Speed", alpha=0.7)
    plt.xlabel("Time Steps")
    plt.ylabel("Lateral Speed")
    plt.legend()
    plt.title("Comparison of Predicted and True Lateral Speeds")
    plt.show()

if __name__ == '__main__':
    test_model()
